package com.william.jdbcplus.core.util.druid;

import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLSetStatement;
import com.alibaba.druid.sql.dialect.db2.visitor.DB2OutputVisitor;
import com.alibaba.druid.sql.dialect.h2.visitor.H2OutputVisitor;
import com.alibaba.druid.sql.dialect.hive.visitor.HiveOutputVisitor;
import com.alibaba.druid.sql.dialect.odps.visitor.OdpsOutputVisitor;
import com.alibaba.druid.sql.dialect.postgresql.visitor.PGOutputVisitor;
import com.alibaba.druid.sql.parser.ParserException;
import com.alibaba.druid.sql.parser.SQLParserFeature;
import com.alibaba.druid.sql.parser.SQLParserUtils;
import com.alibaba.druid.sql.parser.SQLStatementParser;
import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor;
import com.alibaba.druid.sql.visitor.VisitorFeature;
import com.alibaba.druid.support.logging.Log;
import com.alibaba.druid.support.logging.LogFactory;
import com.alibaba.druid.util.JdbcConstants;

import java.lang.reflect.Field;
import java.util.List;
import java.util.Map;

public class SQLUtils extends com.alibaba.druid.sql.SQLUtils {

    private final static SQLParserFeature[] FORMAT_DEFAULT_FEATURES = {SQLParserFeature.KeepComments, SQLParserFeature.EnableSQLBinaryOpExprGroup};

    private final static Log LOG = LogFactory.getLog(SQLUtils.class);

    public static String format(String sql, String dbType, List<Object> parameters, FormatOption option) {
        try {
            SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType, FORMAT_DEFAULT_FEATURES);
            List<SQLStatement> statementList = parser.parseStatementList();
            return toSQLString(statementList, dbType, parameters, option);
        } catch (ClassCastException ex) {
            LOG.warn("format error, dbType : " + dbType, ex);
            return sql;
        } catch (ParserException ex) {
            LOG.warn("format error", ex);
            return sql;
        }
    }

    public static String toSQLString(List<SQLStatement> statementList, String dbType, List<Object> parameters, FormatOption option) {
        return toSQLString(statementList, dbType, parameters, option, null);
    }

    public static String toSQLString(List<SQLStatement> statementList, String dbType, List<Object> parameters, FormatOption option,
                                     Map<String, String> tableMapping) {
        StringBuilder out = new StringBuilder();
        SQLASTOutputVisitor visitor = createFormatOutputVisitor(out, statementList, dbType);
        if (parameters != null) {
            visitor.setInputParameters(parameters);
        }

        if (option == null) {
            option = DEFAULT_FORMAT_OPTION;
        }
        visitor.setFeatures(getFeatures(option));

        if (tableMapping != null) {
            visitor.setTableMapping(tableMapping);
        }

        boolean printStmtSeperator;
        if (JdbcConstants.SQL_SERVER.equals(dbType)) {
            printStmtSeperator = false;
        } else {
            printStmtSeperator = !JdbcConstants.ORACLE.equals(dbType);
        }

        for (int i = 0, size = statementList.size(); i < size; i++) {
            SQLStatement stmt = statementList.get(i);

            if (i > 0) {
                SQLStatement preStmt = statementList.get(i - 1);
                if (printStmtSeperator && !preStmt.isAfterSemi()) {
                    visitor.print(";");
                }

                List<String> comments = preStmt.getAfterCommentsDirect();
                if (comments != null) {
                    for (int j = 0; j < comments.size(); ++j) {
                        String comment = comments.get(j);
                        if (j != 0) {
                            visitor.println();
                        }
                        visitor.printComment(comment);
                    }
                }

                if (printStmtSeperator) {
                    visitor.println();
                }

                if (!(stmt instanceof SQLSetStatement)) {
                    visitor.println();
                }
            }
            {
                List<String> comments = stmt.getBeforeCommentsDirect();
                if (comments != null) {
                    for (String comment : comments) {
                        visitor.printComment(comment);
                        visitor.println();
                    }
                }
            }
            stmt.accept(visitor);

            if (i == size - 1) {
                List<String> comments = stmt.getAfterCommentsDirect();
                if (comments != null) {
                    for (int j = 0; j < comments.size(); ++j) {
                        String comment = comments.get(j);
                        if (j != 0) {
                            visitor.println();
                        }
                        visitor.printComment(comment);
                    }
                }
            }
        }

        return out.toString();
    }

    public static SQLASTOutputVisitor createFormatOutputVisitor(Appendable out, //
                                                                List<SQLStatement> statementList, //
                                                                String dbType) {
        if (JdbcConstants.ORACLE.equals(dbType) || JdbcConstants.ALI_ORACLE.equals(dbType)) {
            if (statementList == null || statementList.size() == 1) {
                return new OracleOutputVisitor(out, false);
            } else {
                return new OracleOutputVisitor(out, true);
            }
        }

        if (JdbcConstants.MYSQL.equals(dbType) //
                || JdbcConstants.MARIADB.equals(dbType)) {
            return new MySqlOutputVisitor(out);
        }

        if (JdbcConstants.POSTGRESQL.equals(dbType)) {
            return new PGOutputVisitor(out);
        }

        if (JdbcConstants.SQL_SERVER.equals(dbType) || JdbcConstants.JTDS.equals(dbType)) {
            return new SQLServerOutputVisitor(out);
        }

        if (JdbcConstants.DB2.equals(dbType)) {
            return new DB2OutputVisitor(out);
        }

        if (JdbcConstants.ODPS.equals(dbType)) {
            return new OdpsOutputVisitor(out);
        }

        if (JdbcConstants.H2.equals(dbType)) {
            return new H2OutputVisitor(out);
        }

        if (JdbcConstants.HIVE.equals(dbType)) {
            return new HiveOutputVisitor(out);
        }

        if (JdbcConstants.ELASTIC_SEARCH.equals(dbType)) {
            return new MySqlOutputVisitor(out);
        }

        return new SQLASTOutputVisitor(out);
    }

    public static int getFeatures(FormatOption option) {
        int fieldValue = VisitorFeature.of(VisitorFeature.OutputUCase, VisitorFeature.OutputPrettyFormat);
        try {
            Field privateField = FormatOption.class.getDeclaredField("features");
            privateField.setAccessible(true);
            fieldValue = (int) privateField.get(option);
        } catch (Exception e) {
            e.printStackTrace();
        }

        return fieldValue;
    }
}
